// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cantor/parser/ast/func_expr.h>
#include <cantor/parser/ast/dql.h>
#include <turbo/flags/flag.h>

TURBO_FLAG(bool, normalize_case_when_to_stream, false,
           "whether reduce case_when/case_expr_when func's children to 2 when to_stream");

namespace cantor::parser {

    std::unordered_map<int, std::string> FuncExpr::FUNC_STR_MAP = {
        {pb::ast::FT_COMMON, ""},
        {pb::ast::FT_AGG, ""},
        // ~ ! -
        {pb::ast::FT_BIT_NOT, "~"},
        {pb::ast::FT_LOGIC_NOT, "!"},
        {pb::ast::FT_UMINUS, "-"},
        // + - * /
        {pb::ast::FT_ADD, " + "},
        {pb::ast::FT_MINUS, " - "},
        {pb::ast::FT_MULTIPLIES, " * "},
        {pb::ast::FT_DIVIDES, " / "},
        // % << >> & | ^
        {pb::ast::FT_MOD, " % "},
        {pb::ast::FT_LS, " << "},
        {pb::ast::FT_RS, " >> "},
        {pb::ast::FT_BIT_AND, " & "},
        {pb::ast::FT_BIT_OR, " | "},
        {pb::ast::FT_BIT_XOR, " ^ "},
        // == != > >= < <=
        {pb::ast::FT_EQ, " = "},
        {pb::ast::FT_NE, " != "},
        {pb::ast::FT_GT, " > "},
        {pb::ast::FT_GE, " >= "},
        {pb::ast::FT_LT, " < "},
        {pb::ast::FT_LE, " <= "},
        // && || xor
        {pb::ast::FT_LOGIC_AND, " && "},
        {pb::ast::FT_LOGIC_OR, " || "},
        {pb::ast::FT_LOGIC_XOR, " XOR "},
        // is like in null ture
        {pb::ast::FT_IS_NULL, "IS NULL"},
        {pb::ast::FT_IS_TRUE, "IS TRUE"},
        {pb::ast::FT_IS_UNKNOWN, "IS UNKNOWN"},
        {pb::ast::FT_IN, "IN"},
        {pb::ast::FT_LIKE, "LIKE"},
        {pb::ast::FT_EXACT_LIKE, "EXACT_LIKE"},
        {pb::ast::FT_REGEXP, "REGEXP"},
        {pb::ast::FT_BETWEEN, "BETWEEN"},
        {pb::ast::FT_VALUES, "VALUES"}
    };

    std::unordered_map<int, const char *> FuncExpr::FUNC_FN_NAME_MAP = {
        {pb::ast::FT_COMMON, ""},
        {pb::ast::FT_AGG, ""},
        // ~ ! -
        {pb::ast::FT_BIT_NOT, "bit_not"},
        {pb::ast::FT_LOGIC_NOT, "logic_not"},
        {pb::ast::FT_UMINUS, "minus"},
        // + - * /
        {pb::ast::FT_ADD, "add"},
        {pb::ast::FT_MINUS, "minus"},
        {pb::ast::FT_MULTIPLIES, "multiplies"},
        {pb::ast::FT_DIVIDES, "divides"},
        // % << >> & | ^
        {pb::ast::FT_MOD, "mod"},
        {pb::ast::FT_LS, "left_shift"},
        {pb::ast::FT_RS, "right_shift"},
        {pb::ast::FT_BIT_AND, "bit_and"},
        {pb::ast::FT_BIT_OR, "bit_or"},
        {pb::ast::FT_BIT_XOR, "bit_xor"},
        // == != > >= < <=
        {pb::ast::FT_EQ, "eq"},
        {pb::ast::FT_NE, "ne"},
        {pb::ast::FT_GT, "gt"},
        {pb::ast::FT_GE, "ge"},
        {pb::ast::FT_LT, "lt"},
        {pb::ast::FT_LE, "le"},
        // && || xor
        {pb::ast::FT_LOGIC_AND, "logic_and"},
        {pb::ast::FT_LOGIC_OR, "logic_or"},
        {pb::ast::FT_LOGIC_XOR, "logic_xor"},
        // is like in null ture
        {pb::ast::FT_IS_NULL, "is_null"},
        {pb::ast::FT_IS_TRUE, "is_true"},
        {pb::ast::FT_IS_UNKNOWN, "is_unknown"},
        {pb::ast::FT_IN, "in"},
        {pb::ast::FT_LIKE, "like"},
        {pb::ast::FT_EXACT_LIKE, "exact_like"},
        {pb::ast::FT_REGEXP, "regexp"},
        {pb::ast::FT_MATCH_AGAINST, "match_against"},
        {pb::ast::FT_BETWEEN, "between"}
    };

    bool FuncExpr::has_subquery() const {
        for (int i = 0; i < exprs.size(); i++) {
            DispatchExpr *expr_node = static_cast<DispatchExpr *>(exprs[i]);
            if (expr_node->is_subquery()) {
                return true;
                }
        }
        return false;
    }

    FuncExpr *FuncExpr::new_unary_op_node(pb::ast::FuncType t, DispatchExpr *arg1, turbo::Arena &arena) {
        FuncExpr *fun = new(arena.allocate(sizeof(FuncExpr)))FuncExpr();
        fun->func_type = t;
        fun->fn_name = type_to_name(t);
        fun->exprs.push_back(arg1, arena);
        return fun;
    }

    FuncExpr *FuncExpr::new_binary_op_node(pb::ast::FuncType t, DispatchExpr *arg1, DispatchExpr *arg2, turbo::Arena &arena) {
        FuncExpr *fun = new(arena.allocate(sizeof(FuncExpr)))FuncExpr();
        fun->func_type = t;
        fun->fn_name = type_to_name(t);
        fun->exprs.push_back(arg1, arena);
        fun->exprs.push_back(arg2, arena);
        return fun;
    }

    FuncExpr *FuncExpr::new_ternary_op_node(
        pb::ast::FuncType t, DispatchExpr *arg1, DispatchExpr *arg2, DispatchExpr *arg3, turbo::Arena &arena) {
        FuncExpr *fun = new(arena.allocate(sizeof(FuncExpr)))FuncExpr();
        fun->func_type = t;
        fun->fn_name = type_to_name(t);
        fun->exprs.reserve(3, arena);
        fun->exprs.push_back(arg1, arena);
        fun->exprs.push_back(arg2, arena);
        fun->exprs.push_back(arg3, arena);
        return fun;
    }

    void FuncExpr::to_stream(std::ostream &os) const {
        static const char *not_str[] = {"", " NOT"};
        static const char *true_str[] = {"TRUE", "FALSE"};
        // Special handling for unary: no parentheses added
        if (indirections.size() > 0) {
            os<<"(";
        }
        switch (func_type) {
            case pb::ast::FT_BIT_NOT:
            case pb::ast::FT_LOGIC_NOT:
            case pb::ast::FT_UMINUS:
                /// SubqueryExpr
                if (exprs[0]->expr_type == pb::ast::ENT_SUB_QUERY_EXPR) {
                    os << FUNC_STR_MAP[func_type] << exprs[0];
                } else {
                    os << "NOT " << exprs[0];
                }
                return;
            default:
                break;
        }
        os << "(";
        switch (func_type) {
            case pb::ast::FT_COMMON:
            case pb::ast::FT_AGG: {
                os << fn_name << "(";
                if (distinct) {
                    os << "DISTINCT ";
                }
                int children_size = exprs.size();
                /// When normalizing functions related to hll/rolling bitmap/tdigest, avoid excessive length.
                if (turbo::get_flag(FLAGS_ast_print_sample) && exprs.size() > 2 &&
                    (fn_name.starts_with("hll_") || fn_name.starts_with("rb_") || fn_name.starts_with("tdigest_")
                     || (turbo::get_flag(FLAGS_normalize_case_when_to_stream) &&
                         (fn_name.to_string() == "case_expr_when" || fn_name.to_string() == "case_when")))) {
                    children_size = 2;
                }
                for (int i = 0; i < children_size; i++) {
                    os << exprs[i];
                    if (i != children_size - 1) {
                        os << ",";
                    }
                }
                if (is_star) {
                    os << "*";
                }
                os << ")";
                break;
            }
            case pb::ast::FT_ADD:
            case pb::ast::FT_MINUS:
            case pb::ast::FT_MULTIPLIES:
            case pb::ast::FT_DIVIDES:
            case pb::ast::FT_MOD:
            case pb::ast::FT_LS:
            case pb::ast::FT_RS:
            case pb::ast::FT_BIT_AND:
            case pb::ast::FT_BIT_OR:
            case pb::ast::FT_BIT_XOR:
            case pb::ast::FT_EQ:
            case pb::ast::FT_NE:
            case pb::ast::FT_GT:
            case pb::ast::FT_GE:
            case pb::ast::FT_LT:
            case pb::ast::FT_LE:
            case pb::ast::FT_LOGIC_AND:
            case pb::ast::FT_LOGIC_OR:
            case pb::ast::FT_LOGIC_XOR:
                os << exprs[0] << FUNC_STR_MAP[func_type] << exprs[1];
                break;
            case pb::ast::FT_IS_NULL:
                os << exprs[0] << " IS" << not_str[is_not] << " NULL";
                break;
            case pb::ast::FT_IS_TRUE:
                os << exprs[0] << " IS " << true_str[is_not];
                break;
            case pb::ast::FT_IS_UNKNOWN:
                os << exprs[0] << " IS" << not_str[is_not] << " UNKNOWN";
                break;
            case pb::ast::FT_IN:
                if (turbo::get_flag(FLAGS_ast_print_sample) && !has_subquery()) {
                    os << exprs[0] << not_str[is_not] << " IN (?)";
                } else {
                    os << exprs[0] << not_str[is_not] << " IN " << exprs[1];
                }
                break;
            case pb::ast::FT_LIKE:
                os << exprs[0] << not_str[is_not] << " LIKE " << exprs[1];
                if (exprs.size() > 2) {
                    os << exprs[2];
                }
                break;
            case pb::ast::FT_EXACT_LIKE:
                os << exprs[0] << not_str[is_not] << " EXACT_LIKE " << exprs[1];
                if (exprs.size() > 2) {
                    os << exprs[2];
                }
                break;
            case pb::ast::FT_REGEXP:
                os << exprs[0] << not_str[is_not] << " REGEXP " << exprs[1];
                break;
            case pb::ast::FT_MATCH_AGAINST:
                os << "MATCH" << exprs[0] << " AGAINST " << "(" << exprs[1] << " " << exprs[2] << ")";
                break;
            case pb::ast::FT_BETWEEN:
                os << exprs[0] << not_str[is_not] <<
                        " BETWEEN " << exprs[1] << " AND " << exprs[2];
                break;
            case pb::ast::FT_VALUES:
                os << "VALUES(" << exprs[0] << ")";
            default:
                break;
        }
        os << ")";
        if (indirections.size() > 0) {
            os<<")";
            for (int i = 0; i < indirections.size(); i++) {
                os << indirections[i];
            }
        }
    }

    nlohmann::ordered_json FuncExpr::to_json() const {
        nlohmann::ordered_json j  = DispatchExpr::to_json_impl();
        j["fn_name"] = fn_name.to_string();
        j["func_type"] = static_cast<int>(func_type);
        j["is_not"] = is_not;
        j["is_star"] = is_star;
        j["distinct"] = distinct;
        auto json_array = nlohmann::ordered_json::array();
        for (auto i = 0 ; i < exprs.size(); i++) {
            json_array.push_back(exprs[i]->to_json());
        }
        j["exprs"] = json_array;
        auto inds = nlohmann::ordered_json::array();
        for (auto i = 0 ; i < indirections.size(); i++) {
            inds.push_back(indirections[i]->to_json());
        }
        j["indirections"] = inds;
        return j;
    }

}  // namespace cantor::parser
